import rclpy
from rclpy.node import Node
from rclpy.clock import Clock
from motor_control_interfaces.msg import MotorsCmd, MotorCmd, MotorsData, MotorData
import asyncio
import serial_asyncio
import struct
import threading

def Emm_V5_Read_Sys_Params(addr, func_codes): # 读取驱动板参数
    i = 0
    cmd = bytearray(16)
    cmd[i] = addr
    i += 1 
    cmd[i] = func_codes
    i += 1
    cmd[i] = 0x6B
    i += 1
    return cmd[:i]

def Pos_Control(addr, vel, clk): # 地址电机，设置方向为CW，速度为1000RPM，加速度为50，脉冲数为2000，相对运动，无多机同步
    cmd = bytearray(13)
    cmd[0] = addr                      # 地址
    cmd[1] = 0xFD                      # 功能码
    if clk<0:
        cmd[2] = 1
        clk=-clk
    else:
        cmd[2] = 0                       # 方向
    cmd[3] = (vel >> 8) & 0xFF         # 速度(RPM)高8位字节
    cmd[4] = vel & 0xFF                # 速度(RPM)低8位字节 
    cmd[5] = 0                       # 加速度，注意：0是直接启动
    cmd[6] = (clk >> 24) & 0xFF        # 脉冲数高8位字节(bit24 - bit31)
    cmd[7] = (clk >> 16) & 0xFF        # 脉冲数(bit16 - bit23)
    cmd[8] = (clk >> 8) & 0xFF         # 脉冲数(bit8  - bit15)
    cmd[9] = clk & 0xFF                # 脉冲数低8位字节(bit0  - bit7)
    cmd[10] = 0x01  # 相位/绝对标志，true为0x01，false为0x00
    cmd[11] = 0x00    # 多机同步运动标志，true为0x01，false为0x00
    cmd[12] = 0x6B                     # 校验字节
    return cmd

def Vel_Control(addr, vel): # 地址电机，设置方向为CW，速度为1000RPM，加速度为50，脉冲数为2000，相对运动，无多机同步
    cmd = bytearray(8)
    cmd[0] = addr                      # 地址
    cmd[1] = 0xF6                      # 功能码
    if vel<0:
        cmd[2] = 1
        vel=-vel
    else:
        cmd[2] = 0                       # 方向
    cmd[3] = (vel >> 8) & 0xFF         # 速度(RPM)高8位字节
    cmd[4] = vel & 0xFF                # 速度(RPM)低8位字节 
    cmd[5] = 0                       # 加速度，注意：0是直接启动
    cmd[6] = 0                      # 00表示不启用多机同步（01表示启用）
    cmd[7] = 0x6B                     # 校验字节
    return cmd


class MotorController(Node):
    def __init__(self):
        super().__init__('motor_controller')
        # 电机数据发布
        self.data_publisher_ = self.create_publisher(MotorsData, 'motors_data', 10)
        self.cmd_subscription_ = self.create_subscription(MotorsCmd, 'motors_cmd', self.cmd_callback, 10)
        self.cmd_subscription_
        self.motors_data = MotorsData()
        self.reader, self.writer = None, None
        # 更新数据
        self.motors_cmd=None
        self.thread = threading.Thread(target=self.run_async_loop, daemon=True)
        self.thread.start()
        

    def cmd_callback(self, msg):
        self.motors_cmd=msg

    def run_async_loop(self):
        asyncio.run(self.StartMotor())

    async def MotorController(self,addr):
        msg = MotorData()
        # 电机实际位置读取
        self.writer.write(Emm_V5_Read_Sys_Params(addr, 0x36))
        data = await self.reader.read(10)
        pos = struct.unpack('>I', data[3:7])[0]
        Motor_Cur_Pos = float(pos) * 360.0 / 65536.0
        if int(data[2]):
            Motor_Cur_Pos = -Motor_Cur_Pos
        msg.position = Motor_Cur_Pos

        # 电机实际速度读取
        self.writer.write(Emm_V5_Read_Sys_Params(addr, 0x35))
        data = await self.reader.read(10)
        vel = struct.unpack('>H', data[3:5])[0]
        if int(data[2]):
            vel = -vel
        msg.velocity = vel

        # 电机实际电流读取
        self.writer.write(Emm_V5_Read_Sys_Params(addr, 0x27))
        data = await self.reader.read(10)
        current = struct.unpack('>H', data[2:4])[0]
        msg.current = current

        # 电机状态
        self.writer.write(Emm_V5_Read_Sys_Params(addr, 0x3A))
        state=await self.reader.read(10)
        msg.enable=bool(state[2] & 0x01)
        msg.stall=bool(state[2] & 0x04)
        self.motors_data.motors_data.append(msg)

        # 电机控制
        if self.motors_cmd is not None:
            # 使能控制
            # 解除堵转控制
            match self.motors_cmd.motors_cmd[addr-1].motor_type:
                # 闭环位置控制
                case 0:
                    self.writer.write(Pos_Control(addr, self.motors_cmd.motors_cmd[addr-1].set_velocity, self.motors_cmd.motors_cmd[addr-1].set_position))
                    state=await self.reader.read(10)
                # 闭环速度控制
                case 1:
                    self.writer.write(Vel_Control(addr, self.motors_cmd.motors_cmd[addr-1].set_velocity))
                    state=await self.reader.read(10)
                case _:
                    print('未知电机类型')
        
    async def StartMotor(self):
        self.reader, self.writer = await serial_asyncio.open_serial_connection(url='/dev/ttyUSB0', baudrate=115200)
        while rclpy.ok():      
            self.motors_data = MotorsData()
            # tasks = [self.MotorController(i) for i in range(1, 3)]  # 生成任务列表
            # await asyncio.gather(*tasks)  # 并发执行所有电机数据读取
            for i in range(1, 3):  # 顺序执行所有电机数据读取
                await self.MotorController(i)
            self.data_publisher_.publish(self.motors_data)


    
def main(args=None):
    rclpy.init(args=args)
    motor_controller = MotorController()
    rclpy.spin(motor_controller)
    motor_controller.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()